import torch
import torch.nn as nn

# 准备数据
batch_size = 8
torch.manual_seed(42)
x = torch.rand(batch_size, 64, 64)
y = torch.tensor([1., 1., 0., 1., 0., 0., 0., 0.]).unsqueeze(1)
# print(y)

# 1. 超参数
lr = 0.01
Epochs = 500

# 2. 定义模型
model = nn.Sequential(
    nn.Flatten(start_dim=1),
    nn.Linear(64*64, 1024),
    nn.ReLU(),
    nn.Linear(1024, 1),
    nn.Sigmoid()
)

'''
最后一层，最开始写的out_features=2，表示的不再是单纯的单标签二分类，而是两个标签的二分类
比如：
狗：田园犬、金毛
猫：橘猫、三花
若 out_features=1，则只判断 是猫还是狗
'''

# 3. 损失函数
# nn.BCEWithLogitsLoss()      自带Sigmoid激活
loss_fn = nn.BCELoss()        # 关于二元交叉熵损失函数 和 交叉熵损失函数的对比

# 4. 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 5. 循环训练
for epoch in range(Epochs):
    # 前向传播
    y_pre = model(x)
    # 计算损失
    loss = loss_fn(y_pre, y)
    # 清空梯度
    optimizer.zero_grad()
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()

    # 打印训练信息
    if epoch == 0 or (epoch + 1) % 10 == 0:
        print(f"[{epoch + 1}/{Epochs}] Loss:{loss.item():.4f}")

y_pre = model(x)
print(f"真实标签：{y.squeeze()}")
print(f"预测标签：{(y_pre.squeeze()>=0.5)}")